{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 将单机模型训练代码迁移 SecretFlow 联邦学习训练代码教程"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 引言\n",
    "### 背景\n",
    "随着数据隐私问题日益受到重视，并且实际业务场景的需求不断演变，联邦学习作为一种特殊的深度学习形式开始迅速兴起。它以一种创新的方式解决了传统集中式训练的数据隐私问题，能够在保护用户隐私的前提下，有效地训练机器学习模型。  \n",
    "那么假如现实的业务中已经有了单机模型，应该如何将它进行联邦化呢？为了易用性，我们的隐语在设计之初就希望能够以最低的成本来帮助用户迁移已有的单机模型到联邦模型。\n",
    "\n",
    "### 教程目标以及内容\n",
    "通过本教程的学习：\n",
    "1. 读者可以快速的将已有的单机模型正确的使用SecretFlow来进行联邦化。  \n",
    "2. 先通过单机模型完成模型开发，再进行联邦化适配，也是一个比较推荐的开发实践，可以提高联邦模型的开发效率。\n",
    "\n",
    "本教程将手把手的带你学习如何将已有的单机模型训练代码迁移到SecretFlow中进行联邦学习。\n",
    "1. 介绍迁移的整体流程\n",
    "2. 通过案例介绍在Tensorflow作为后端的迁移流程\n",
    "3. 通过案例介绍在Pytorch作为后端的迁移流程  \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 迁移的步骤\n",
    "\n",
    "基于 PyTorch 从单机模型到联邦学习模型，主要需要添加或修改以下几部分\n",
    "- 添加联邦学习中的参与方\n",
    "- 修改数据集的处理逻辑\n",
    "- 修改模型的继承类\n",
    "- 根据需要决定是否对 metric 、 optimizer 和 loss fuction 进行包装，并使用包装后的函数\n",
    "\n",
    "基于 TensorFlow 从单机模型到联邦学习模型，主要需要添加或修改以下几部分\n",
    "- 添加联邦学习中的参与方\n",
    "- 修改数据集的处理逻辑\n",
    "- 进行模型的封装\n",
    "\n",
    "得益于隐语的封装，使用者不需要自己进行大量的代码编写，只需要调用 Secretflow 中的函数，即可便捷完成模型的定义和使用等操作。  \n",
    "迁移完成后，不同后端不同的模型都使用一套API来进行`fit`,`predict`,`evaluate`等等"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 前期数据准备\n",
    "相关文档可以参考相关IO文档\n",
    "\n",
    "下载数据集并解压"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import requests\n",
    "import tarfile\n",
    "import tempfile\n",
    "\n",
    "# Create a temporary folder\n",
    "_temp_dir = tempfile.mkdtemp()\n",
    "\n",
    "# Download file\n",
    "url = \"https://secretflow-data.oss-accelerate.aliyuncs.com/datasets/tf_flowers/flower_photos.tgz\"\n",
    "save_path = os.path.join(_temp_dir, \"flower_photos.tgz\")\n",
    "\n",
    "response = requests.get(url)\n",
    "with open(save_path, \"wb\") as f:\n",
    "    f.write(response.content)\n",
    "\n",
    "# Extract the file\n",
    "extract_folder = os.path.join(_temp_dir, \"flower_photos\")\n",
    "os.makedirs(extract_folder, exist_ok=True)\n",
    "\n",
    "with tarfile.open(save_path, \"r:gz\") as tar:\n",
    "    tar.extractall(path=extract_folder)\n",
    "\n",
    "path_to_flower_dataset = extract_folder"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 基于PyTorch的迁移教程\n",
    "\n",
    "### 模型在 PyTorch 下的单机模型实现\n",
    "首先我们给出单机模式下，基于 PyTorch 定义和训练神经网络模型的过程。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "\n",
    "import numpy as np\n",
    "from torch.utils.data import DataLoader\n",
    "from torch.utils.data.sampler import SubsetRandomSampler\n",
    "from torchvision import datasets, transforms\n",
    "\n",
    "# parameter\n",
    "batch_size = 32\n",
    "shuffle = True\n",
    "random_seed = 1234\n",
    "train_split = 0.8\n",
    "\n",
    "# 构建 PyTorch 中 Dataloader 对象\n",
    "flower_transform = transforms.Compose(\n",
    "    [\n",
    "        transforms.Resize((180, 180)),\n",
    "        transforms.ToTensor(),\n",
    "    ]\n",
    ")\n",
    "flower_dataset = datasets.ImageFolder(\n",
    "    path_to_flower_dataset, transform=flower_transform\n",
    ")\n",
    "dataset_size = len(flower_dataset)\n",
    "# Define sampler\n",
    "\n",
    "indices = list(range(dataset_size))\n",
    "if shuffle:\n",
    "    np.random.seed(random_seed)\n",
    "    np.random.shuffle(indices)\n",
    "split = int(np.floor(train_split * dataset_size))\n",
    "train_indices, val_indices = indices[:split], indices[split:]\n",
    "train_sampler = SubsetRandomSampler(train_indices)\n",
    "valid_sampler = SubsetRandomSampler(val_indices)\n",
    "\n",
    "# Define databuilder\n",
    "train_loader = DataLoader(flower_dataset, batch_size=batch_size, sampler=train_sampler)\n",
    "valid_loader = DataLoader(flower_dataset, batch_size=batch_size, sampler=valid_sampler)\n",
    "\n",
    "\n",
    "# 定义单机模型结构\n",
    "import torch\n",
    "from torch import nn\n",
    "\n",
    "\n",
    "class ConvRGBNet_torch(nn.Module):\n",
    "    def __init__(self, *args, **kwargs) -> None:\n",
    "        super().__init__(*args, **kwargs)\n",
    "        self.network = nn.Sequential(\n",
    "            nn.Conv2d(\n",
    "                in_channels=3, out_channels=16, kernel_size=3, stride=1, padding=1\n",
    "            ),\n",
    "            nn.ReLU(),\n",
    "            nn.MaxPool2d(2, 2),\n",
    "            nn.Conv2d(16, 16, kernel_size=3, stride=1, padding=1),\n",
    "            nn.ReLU(),\n",
    "            nn.MaxPool2d(2, 2),\n",
    "            nn.Flatten(),\n",
    "            nn.Linear(16 * 45 * 45, 128),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(128, 5),\n",
    "        )\n",
    "\n",
    "    def forward(self, xb):\n",
    "        return self.network(xb)\n",
    "\n",
    "\n",
    "# 定义优化器\n",
    "# initialize\n",
    "model_torch = ConvRGBNet_torch()\n",
    "\n",
    "# Define the loss function\n",
    "loss_model_torch = nn.CrossEntropyLoss()\n",
    "\n",
    "# Define the optimizer\n",
    "optimizer_model_torch = torch.optim.SGD(params=model_torch.parameters(), lr=0.01)\n",
    "optimizer_model_torch.zero_grad()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 基于 PyTorch 的隐语联邦学习模型迁移\n",
    "基于 PyTorch 从单机模型到联邦学习模型，主要包含以下步骤：\n",
    "- 添加联邦学习中的参与方\n",
    "- 修改数据集的处理逻辑\n",
    "- 修改模型的继承类\n",
    "- 根据需要决定是否对 metric 、 optimizer 和 loss fuction 进行包装，并使用包装后的函数\n",
    "\n",
    "接下来，我们将结合实际代码具体讲解这些步骤。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 环境设置\n",
    "添加联邦学习中的参与方，并对各个参与方进行初始化"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The version of SecretFlow: 1.5.0.dev20240304\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/user/miniconda3/envs/secretflow/lib/python3.10/subprocess.py:1796: RuntimeWarning: os.fork() was called. os.fork() is incompatible with multithreaded code, and JAX is multithreaded, so this will likely lead to a deadlock.\n",
      "  self.pid = _posixsubprocess.fork_exec(\n",
      "2024-03-07 17:17:40,767\tINFO worker.py:1724 -- Started a local Ray instance.\n"
     ]
    }
   ],
   "source": [
    "import secretflow as sf\n",
    "\n",
    "# Check the version of your SecretFlow\n",
    "print('The version of SecretFlow: {}'.format(sf.__version__))\n",
    "\n",
    "# In case you have a running secretflow runtime already.\n",
    "sf.shutdown()\n",
    "sf.init(['alice', 'bob', 'charlie'], address=\"local\", log_to_driver=False)\n",
    "alice, bob, charlie = sf.PYU('alice'), sf.PYU('bob'), sf.PYU('charlie')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 封装单机模式下的数据处理逻辑\n",
    "同样地，在联邦学习中我们也需要对数据进行预处理，使之符合模型的输入，所以参考在[SecretFlow 中使用自定义 DataBuilder (Torch)构建 dataset builder](https://www.secretflow.org.cn/docs/secretflow/latest/zh-Hans/tutorial/CustomDataLoaderTorch)，我们选择文件夹路径作为参数，并且封装单机模式下的数据处理逻辑，最后返回 (data_set，steps_per_epoch)的结果，封装代码如下："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_dataset_builder(\n",
    "    batch_size=32,\n",
    "    train_split=0.8,\n",
    "    shuffle=True,\n",
    "    random_seed=1234,\n",
    "):\n",
    "    def dataset_builder(x, stage=\"train\"):\n",
    "        \"\"\" \"\"\"\n",
    "        ######################################单机模型数据读取代码#########################################\n",
    "        import math\n",
    "\n",
    "        import numpy as np\n",
    "        from torch.utils.data import DataLoader\n",
    "        from torch.utils.data.sampler import SubsetRandomSampler\n",
    "        from torchvision import datasets, transforms\n",
    "\n",
    "        # Define dataset\n",
    "\n",
    "        flower_transform = transforms.Compose(\n",
    "            [\n",
    "                transforms.Resize((180, 180)),\n",
    "                transforms.ToTensor(),\n",
    "            ]\n",
    "        )\n",
    "        flower_dataset = datasets.ImageFolder(x, transform=flower_transform)\n",
    "        dataset_size = len(flower_dataset)\n",
    "        # Define sampler\n",
    "\n",
    "        indices = list(range(dataset_size))\n",
    "        if shuffle:\n",
    "            np.random.seed(random_seed)\n",
    "            np.random.shuffle(indices)\n",
    "        split = int(np.floor(train_split * dataset_size))\n",
    "        train_indices, val_indices = indices[:split], indices[split:]\n",
    "        train_sampler = SubsetRandomSampler(train_indices)\n",
    "        valid_sampler = SubsetRandomSampler(val_indices)\n",
    "\n",
    "        # Define databuilder\n",
    "        train_loader = DataLoader(\n",
    "            flower_dataset, batch_size=batch_size, sampler=train_sampler\n",
    "        )\n",
    "        valid_loader = DataLoader(\n",
    "            flower_dataset, batch_size=batch_size, sampler=valid_sampler\n",
    "        )\n",
    "        #############################################################################################\n",
    "        # Return\n",
    "        if stage == \"train\":\n",
    "            train_step_per_epoch = len(train_loader)\n",
    "\n",
    "            return train_loader, train_step_per_epoch\n",
    "        elif stage == \"eval\":\n",
    "            eval_step_per_epoch = len(valid_loader)\n",
    "            return valid_loader, eval_step_per_epoch\n",
    "\n",
    "    return dataset_builder"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 构建 dataset_builder_dict\n",
    "\n",
    "我们通过 dataset_builder_dict 为各个参与方传入封装数据处理逻辑的 create_dataset_builder 函数的参数。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# prepare dataset dict\n",
    "data_builder_dict = {\n",
    "    alice: create_dataset_builder(\n",
    "        batch_size=32,\n",
    "        train_split=0.8,\n",
    "        shuffle=False,\n",
    "        random_seed=1234,\n",
    "    ),\n",
    "    bob: create_dataset_builder(\n",
    "        batch_size=32,\n",
    "        train_split=0.8,\n",
    "        shuffle=False,\n",
    "        random_seed=1234,\n",
    "    ),\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 在隐语的框架下定义基于 PyTorch 的模型架构\n",
    "参考 PyTorch 单机模式下的模型，我们在隐语的框架下定义同样结构的模型。\n",
    "\n",
    "我们只需修改继承类将 **torch.nn.Module** 改为 **secretflow.ml.nn.core.torch.BaseModule**，就可以完成模型架构的定义。\n",
    "\n",
    "从迁移过程可以看出，将单机模型在隐语框架下进行定义所进行的代码改动非常小，整体迁移非常方便，充分展现了隐语框架的易用性。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "from secretflow.ml.nn.core.torch import BaseModule\n",
    "\n",
    "\n",
    "######################################单机模型代码#########################################\n",
    "class ConvRGBNet(BaseModule):  # 只有这里有变化，其他和单机模型定义保持一致\n",
    "    def __init__(self, *args, **kwargs) -> None:\n",
    "        super().__init__(*args, **kwargs)\n",
    "        self.network = nn.Sequential(\n",
    "            nn.Conv2d(\n",
    "                in_channels=3, out_channels=16, kernel_size=3, stride=1, padding=1\n",
    "            ),\n",
    "            nn.ReLU(),\n",
    "            nn.MaxPool2d(2, 2),\n",
    "            nn.Conv2d(16, 16, kernel_size=3, stride=1, padding=1),\n",
    "            nn.ReLU(),\n",
    "            nn.MaxPool2d(2, 2),\n",
    "            nn.Flatten(),\n",
    "            nn.Linear(16 * 45 * 45, 128),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(128, 5),\n",
    "        )\n",
    "\n",
    "    def forward(self, xb):\n",
    "        return self.network(xb)\n",
    "\n",
    "\n",
    "################################################################################"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 构建TorchModel\n",
    "`TorchModel`是我们在隐语中定义的一个概念，层级概念上对应`keras model`。目的是将将用户定义的`模型`，`loss函数`,`optimizer`,`metrics`封装在一起，用与后续FLModel在启动后统一多后端逻辑。  \n",
    "其中需要注意的是，我们的loss_fn，optim_fn，metrics都需要通过一个wrapper封装成一个偏函数传入进去，在后面会有详细的解析。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:Create proxy actor <class 'secretflow.security.aggregation.secure_aggregator._Masker'> with party alice.\n",
      "INFO:root:Create proxy actor <class 'secretflow.security.aggregation.secure_aggregator._Masker'> with party bob.\n"
     ]
    }
   ],
   "source": [
    "from secretflow.ml.nn import FLModel\n",
    "from secretflow.security.aggregation import SecureAggregator\n",
    "from torch import nn, optim\n",
    "from torchmetrics import Accuracy, Precision\n",
    "from secretflow.ml.nn.core.torch import TorchModel, metric_wrapper, optim_wrapper\n",
    "\n",
    "\n",
    "device_list = [alice, bob]\n",
    "aggregator = SecureAggregator(charlie, [alice, bob])\n",
    "# prepare model\n",
    "num_classes = 5\n",
    "\n",
    "\n",
    "input_shape = (180, 180, 3)\n",
    "# torch model\n",
    "\n",
    "optim_fn = optim_wrapper(optim.Adam, lr=1e-3)\n",
    "model_def = TorchModel(\n",
    "    model_fn=ConvRGBNet,\n",
    "    loss_fn=nn.CrossEntropyLoss,\n",
    "    optim_fn=optim_fn,\n",
    "    metrics=[\n",
    "        metric_wrapper(\n",
    "            Accuracy, task=\"multiclass\", num_classes=num_classes, average='micro'\n",
    "        ),\n",
    "        metric_wrapper(\n",
    "            Precision, task=\"multiclass\", num_classes=num_classes, average='micro'\n",
    "        ),\n",
    "    ],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 定义FLModel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:Create proxy actor <class 'secretflow.ml.nn.fl.backend.torch.strategy.fed_avg_w.PYUFedAvgW'> with party alice.\n",
      "INFO:root:Create proxy actor <class 'secretflow.ml.nn.fl.backend.torch.strategy.fed_avg_w.PYUFedAvgW'> with party bob.\n"
     ]
    }
   ],
   "source": [
    "fed_model = FLModel(\n",
    "    device_list=device_list,\n",
    "    model=model_def,\n",
    "    aggregator=aggregator,\n",
    "    backend=\"torch\",  # backend support ['tensorflow', 'torch']\n",
    "    strategy=\"fed_avg_w\",\n",
    "    random_seed=1234,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:FL Train Params: {'x': {PYURuntime(alice): '/tmp/tmp32k3ewwr/flower_photos', PYURuntime(bob): '/tmp/tmp32k3ewwr/flower_photos'}, 'y': None, 'batch_size': 32, 'batch_sampling_rate': None, 'epochs': 5, 'verbose': 1, 'callbacks': None, 'validation_data': {PYURuntime(alice): '/tmp/tmp32k3ewwr/flower_photos', PYURuntime(bob): '/tmp/tmp32k3ewwr/flower_photos'}, 'shuffle': False, 'class_weight': None, 'sample_weight': None, 'validation_freq': 1, 'aggregate_freq': 2, 'label_decoder': None, 'max_batch_size': 20000, 'prefetch_buffer_size': None, 'sampler_method': 'batch', 'random_seed': 1234, 'dp_spent_step_freq': 1, 'audit_log_dir': None, 'dataset_builder': {PYURuntime(alice): <function create_dataset_builder.<locals>.dataset_builder at 0x7fd38ac54dc0>, PYURuntime(bob): <function create_dataset_builder.<locals>.dataset_builder at 0x7fd3a412e3b0>}, 'wait_steps': 100, 'self': <secretflow.ml.nn.fl.fl_model.FLModel object at 0x7fd37caf3940>}\n",
      "Train Processing: :   0%|          | 0/30 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/5\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train Processing: :  93%|█████████▎| 28/30 [00:16<00:01,  1.88it/s]/home/user/miniconda3/envs/secretflow/lib/python3.10/site-packages/secretflow/ml/nn/metrics.py:59: UserWarning: Please pay attention to local metrics, global only do naive aggregation.\n",
      "  warnings.warn(\n",
      "2024-03-07 17:18:04.077106: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer.so.7'; dlerror: libnvinfer.so.7: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /usr/local/cuda-11/lib64:\n",
      "2024-03-07 17:18:04.077206: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer_plugin.so.7'; dlerror: libnvinfer_plugin.so.7: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /usr/local/cuda-11/lib64:\n",
      "2024-03-07 17:18:04.077213: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Cannot dlopen some TensorRT libraries. If you would like to use Nvidia GPU with TensorRT, please make sure the missing libraries mentioned above are installed properly.\n",
      "Train Processing: :  93%|█████████▎| 28/30 [00:19<00:01,  1.41it/s, {'multiclassaccuracy': 0.96666664, 'multiclassprecision': 0.96666664, 'val_multiclassaccuracy': 1.0, 'val_multiclassprecision': 1.0}]\n",
      "Train Processing: :   0%|          | 0/30 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 2/5\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train Processing: :  93%|█████████▎| 28/30 [00:17<00:01,  1.58it/s, {'multiclassaccuracy': 1.0, 'multiclassprecision': 1.0, 'val_multiclassaccuracy': 1.0, 'val_multiclassprecision': 1.0}]\n",
      "Train Processing: :   0%|          | 0/30 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 3/5\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train Processing: :  93%|█████████▎| 28/30 [00:17<00:01,  1.57it/s, {'multiclassaccuracy': 1.0, 'multiclassprecision': 1.0, 'val_multiclassaccuracy': 1.0, 'val_multiclassprecision': 1.0}]\n",
      "Train Processing: :   0%|          | 0/30 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 4/5\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train Processing: :  93%|█████████▎| 28/30 [00:18<00:01,  1.55it/s, {'multiclassaccuracy': 1.0, 'multiclassprecision': 1.0, 'val_multiclassaccuracy': 1.0, 'val_multiclassprecision': 1.0}]\n",
      "Train Processing: :   0%|          | 0/30 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 5/5\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train Processing: :  93%|█████████▎| 28/30 [00:18<00:01,  1.50it/s, {'multiclassaccuracy': 1.0, 'multiclassprecision': 1.0, 'val_multiclassaccuracy': 1.0, 'val_multiclassprecision': 1.0}]\n"
     ]
    }
   ],
   "source": [
    "data = {\n",
    "    alice: path_to_flower_dataset,\n",
    "    bob: path_to_flower_dataset,\n",
    "}\n",
    "history = fed_model.fit(\n",
    "    data,\n",
    "    None,\n",
    "    validation_data=data,\n",
    "    epochs=5,\n",
    "    batch_size=32,\n",
    "    aggregate_freq=2,\n",
    "    sampler_method=\"batch\",\n",
    "    random_seed=1234,\n",
    "    dp_spent_step_freq=1,\n",
    "    dataset_builder=data_builder_dict,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 构建完成"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "得到FLModel对象之后，就可以用这个对象进行`fit`,`evaluate`,`predict`等操作了。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Why need wraps"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 将单机模型下的 optimizer 包装（wrap）\n",
    "对 optimizer 进行包装的原因：我们需要使用 optimizer 在训练过程中自动地调整模型参数，以使模型在给定的训练数据上达到最佳的性能表现。同样地，联邦学习也可以通过 optimizer 实现模型更好的性能。但由于隐语的设计机制需要把相关模块通过序列化到具体的机器上才会执行，因此在优化器需要指定参数时，需要做一次封装。\n",
    "\n",
    "通过`optim_wrapper`进行优化器（optimizer）的包装，并且通过追根溯源，我们可以看到。\n",
    "```python\n",
    "from secretflow.ml.nn.core.torch import optim_wrapper\n",
    "```\n",
    "[source code](https://github.com/secretflow/secretflow/blob/main/secretflow/ml/nn/core/torch/utils.py#L23)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def optim_wrapper(func, *args, **kwargs):\n",
    "\n",
    "    def wrapped_func(params):\n",
    "        return func(params, *args, **kwargs)\n",
    "\n",
    "    return wrapped_func"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "可以看到函数实际上都是通过传入一个需要包装的函数名称，位置参数和关键字参数对函数完成包装，通过关键字参数确保了指定的参数赋值，然后返回包装好的函数。\n",
    "因此\n",
    "```python\n",
    "optim_fn = optim_wrapper(optim.Adam, lr=1e-2)\n",
    "```\n",
    "实际上相当于调用\n",
    "```python\n",
    "optim.Adam(lr=1e-2)\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 将单机模型下的 metric 包装（wrap）\n",
    "对 metric 进行包装的原因:metric 在机器学习和深度学习中用于衡量模型的性能和表现，它们是评估模型在训练、验证或测试数据上的效果的标准。同样地，我们也希望使用 metric 衡量联邦学习模型的性能和表现。但由于隐语的设计机制同 PyTorch 有些差异，因此在衡量指标需要指定参数时，需要做一次封装来保证两者的一致性。\n",
    "\n",
    "通过`metric_wrapper`进行衡量指标（metric）的包装，并且通过追根溯源，我们可以看到\n",
    "```python\n",
    "from secretflow.ml.nn.core.torch import metric_wrapper\n",
    "```\n",
    "[source code](https://github.com/secretflow/secretflow/blob/main/secretflow/ml/nn/core/torch/utils.py#L23)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def metric_wrapper(func, *args, **kwargs):\n",
    "    def wrapped_func():\n",
    "        return func(*args, **kwargs)\n",
    "\n",
    "    return wrapped_func"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "可以看到函数实际上是通过传入一个需要包装的函数名称，位置参数和关键字参数对函数完成包装，通过关键字参数确保了指定的参数赋值，然后返回包装好的函数；与优化器包装类似\n",
    "因此\n",
    "```python\n",
    "metric_wrapper(Accuracy, task=\"multiclass\", num_classes=10, average='micro')\n",
    "```\n",
    "实际上相当于调用\n",
    "```python\n",
    "Accuracy(task=\"multiclass\", num_classes=10, average='micro')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 将单机模型下的 loss function 包装（wrap）\n",
    "参考 `optim_wrapper` 和 `metric_wrapper` 的定义方式，自定义 `loss_function_wrapper`对损失函数（loss function）进行包装。\n",
    "\n",
    "如果模型使用损失函数的默认参数，则不需要使用包装"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "def loss_function_wrapper(func, *args, **kwargs):\n",
    "    def wrapped_func():\n",
    "        return func(*args, **kwargs)\n",
    "\n",
    "    return wrapped_func"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "得益于隐语的封装，并且根据前述的包装原理可知，我们实际上需要通过包装完成一个参数具体化的函数，所以在这里 `nn.CrossEntropyLoss` 是需要包装的函数，并且其参数取值，就是需要传入的参数值。因为在其默认参数取值设置中，  `reduction='mean'` , 此处我们试着将其修改为 ` reduction='sum'`。包装损失函数只需要写成："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "loss_wrapper = loss_function_wrapper(nn.CrossEntropyLoss, reduction='sum')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 对单机模型使用包装（wrap）后的优化器（optimizer）、衡量指标（metric）和损失函数（loss function）\n",
    "在隐语框架下，使用者可以根据需要选择是否对优化器（optimizer）、衡量指标（metric）和损失函数（loss function）进行包装，从而更好地训练联邦学习模型，充分发挥隐语框架的灵活性。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 小结\n",
    "基于 PyTorch 从单机模型到联邦学习模型，主要需要添加或修改以下几部分\n",
    "- 添加联邦学习中的参与方\n",
    "- 修改数据集的处理逻辑\n",
    "- 修改模型的继承类\n",
    "- 根据需要决定是否对 metric 、 optimizer 和 loss fuction 进行包装，并使用包装后的函数\n",
    "\n",
    "得益于隐语的封装，使用者不需要自己完成模型定义等代码的编写，只需要调用 Secretflow 中的函数，即可便捷完成模型的定义和使用等操作。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 基于 TensorFlow 的迁移教程\n",
    "### 模型在 TensorFlow 下的单机模型实现\n",
    "首先我们给出单机模式下，基于 TensorFlow 定义和训练神经网络模型的过程。对于数据，我们将其加载成 TensorFlow 的 dataset 对象\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Found 1201 files belonging to 1 classes.\n",
      "Using 961 files for training.\n",
      "Using 240 files for validation.\n",
      "-------Start training-------\n",
      "Epoch 1/10\n",
      "31/31 [==============================] - 2s 17ms/step - loss: 0.0601 - accuracy: 0.9667 - val_loss: 0.0000e+00 - val_accuracy: 1.0000\n",
      "Epoch 2/10\n",
      "31/31 [==============================] - 0s 11ms/step - loss: 0.0000e+00 - accuracy: 1.0000 - val_loss: 0.0000e+00 - val_accuracy: 1.0000\n",
      "Epoch 3/10\n",
      "31/31 [==============================] - 0s 9ms/step - loss: 0.0000e+00 - accuracy: 1.0000 - val_loss: 0.0000e+00 - val_accuracy: 1.0000\n",
      "Epoch 4/10\n",
      "31/31 [==============================] - 0s 10ms/step - loss: 0.0000e+00 - accuracy: 1.0000 - val_loss: 0.0000e+00 - val_accuracy: 1.0000\n",
      "Epoch 5/10\n",
      "31/31 [==============================] - 0s 9ms/step - loss: 0.0000e+00 - accuracy: 1.0000 - val_loss: 0.0000e+00 - val_accuracy: 1.0000\n",
      "Epoch 6/10\n",
      "31/31 [==============================] - 0s 10ms/step - loss: 0.0000e+00 - accuracy: 1.0000 - val_loss: 0.0000e+00 - val_accuracy: 1.0000\n",
      "Epoch 7/10\n",
      "31/31 [==============================] - 0s 10ms/step - loss: 0.0000e+00 - accuracy: 1.0000 - val_loss: 0.0000e+00 - val_accuracy: 1.0000\n",
      "Epoch 8/10\n",
      "31/31 [==============================] - 0s 9ms/step - loss: 0.0000e+00 - accuracy: 1.0000 - val_loss: 0.0000e+00 - val_accuracy: 1.0000\n",
      "Epoch 9/10\n",
      "31/31 [==============================] - 0s 9ms/step - loss: 0.0000e+00 - accuracy: 1.0000 - val_loss: 0.0000e+00 - val_accuracy: 1.0000\n",
      "Epoch 10/10\n",
      "31/31 [==============================] - 0s 9ms/step - loss: 0.0000e+00 - accuracy: 1.0000 - val_loss: 0.0000e+00 - val_accuracy: 1.0000\n"
     ]
    }
   ],
   "source": [
    "import math\n",
    "import tensorflow as tf\n",
    "\n",
    "img_height = 180\n",
    "img_width = 180\n",
    "batch_size = 32\n",
    "# In this example, we use the TensorFlow interface for development.\n",
    "data_set = tf.keras.utils.image_dataset_from_directory(\n",
    "    path_to_flower_dataset,\n",
    "    validation_split=0.2,\n",
    "    subset=\"both\",\n",
    "    seed=123,\n",
    "    image_size=(img_height, img_width),\n",
    "    batch_size=batch_size,\n",
    ")\n",
    "train_set = data_set[0]\n",
    "test_set = data_set[1]\n",
    "\n",
    "from tensorflow import keras\n",
    "\n",
    "# Create model\n",
    "num_classes = 5\n",
    "input_shape = (180, 180, 3)\n",
    "total_epochs = 10\n",
    "\n",
    "model_tensorflow = keras.Sequential(\n",
    "    [\n",
    "        keras.Input(shape=input_shape),\n",
    "        tf.keras.layers.Rescaling(1.0 / 255),\n",
    "        tf.keras.layers.Conv2D(32, 3, activation='relu'),\n",
    "        tf.keras.layers.MaxPooling2D(),\n",
    "        tf.keras.layers.Conv2D(32, 3, activation='relu'),\n",
    "        tf.keras.layers.MaxPooling2D(),\n",
    "        tf.keras.layers.Conv2D(32, 3, activation='relu'),\n",
    "        tf.keras.layers.MaxPooling2D(),\n",
    "        tf.keras.layers.Flatten(),\n",
    "        tf.keras.layers.Dense(128, activation='relu'),\n",
    "        tf.keras.layers.Dense(num_classes),\n",
    "    ]\n",
    ")\n",
    "\n",
    "# Compile model\n",
    "model_tensorflow.compile(\n",
    "    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
    "    optimizer='adam',\n",
    "    metrics=[\"accuracy\"],\n",
    ")\n",
    "\n",
    "# Model training and validation\n",
    "print('-------Start training-------')\n",
    "history = model_tensorflow.fit(\n",
    "    train_set,\n",
    "    validation_data=test_set,\n",
    "    batch_size=batch_size,\n",
    "    epochs=total_epochs,\n",
    "    verbose=True,\n",
    "    shuffle=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 基于 TensorFlow 的隐语联邦学习模型迁移\n",
    "#### 概述 \n",
    "基于 TensorFlow 从单机模型到联邦学习模型，主要包含以下步骤：\n",
    "- 添加联邦学习中的参与方\n",
    "- 修改数据集的处理逻辑\n",
    "- 进行模型的封装\n",
    "\n",
    "接下来，我们将结合实际代码具体讲解这些步骤。\n",
    "### 环境设置\n",
    "添加联邦学习中的参与方，并初始化各个参与方"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The autoreload extension is already loaded. To reload it, use:\n",
      "  %reload_ext autoreload\n"
     ]
    }
   ],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The version of SecretFlow: 1.5.0.dev20240304\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-03-07 17:19:26,589\tINFO worker.py:1724 -- Started a local Ray instance.\n"
     ]
    }
   ],
   "source": [
    "import secretflow as sf\n",
    "\n",
    "# Check the version of your SecretFlow\n",
    "print('The version of SecretFlow: {}'.format(sf.__version__))\n",
    "\n",
    "# In case you have a running secretflow runtime already.\n",
    "sf.shutdown()\n",
    "sf.init(['alice', 'bob', 'charlie'], address=\"local\", log_to_driver=False)\n",
    "alice, bob, charlie = sf.PYU('alice'), sf.PYU('bob'), sf.PYU('charlie')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_dataset_builder(\n",
    "    batch_size=32,\n",
    "):\n",
    "    def dataset_builder(folder_path, stage=\"train\"):\n",
    "        ############################单机模型代码#############################################\n",
    "        import math\n",
    "\n",
    "        import tensorflow as tf\n",
    "\n",
    "        img_height = 180\n",
    "        img_width = 180\n",
    "        data_set = tf.keras.utils.image_dataset_from_directory(\n",
    "            folder_path,\n",
    "            validation_split=0.2,\n",
    "            subset=\"both\",\n",
    "            seed=123,\n",
    "            image_size=(img_height, img_width),\n",
    "            batch_size=batch_size,\n",
    "        )\n",
    "        ####################################################################################\n",
    "        if stage == \"train\":\n",
    "            train_dataset = data_set[0]\n",
    "            train_step_per_epoch = math.ceil(len(data_set[0].file_paths) / batch_size)\n",
    "            return train_dataset, train_step_per_epoch\n",
    "        elif stage == \"eval\":\n",
    "            eval_dataset = data_set[1]\n",
    "            eval_step_per_epoch = math.ceil(len(data_set[1].file_paths) / batch_size)\n",
    "            return eval_dataset, eval_step_per_epoch\n",
    "\n",
    "    return dataset_builder"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 构建 dataset_builder_dict\n",
    "我们通过 dataset_builder_dict 为各个参与方传入封装数据处理逻辑的 create_dataset_builder 函数的参数。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_builder_dict = {\n",
    "    alice: create_dataset_builder(\n",
    "        batch_size=32,\n",
    "    ),\n",
    "    bob: create_dataset_builder(\n",
    "        batch_size=32,\n",
    "    ),\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 在隐语的框架下定义基于 TensorFlow 的模型架构\n",
    "参考 TensorFlow 单机模式下的模型，我们在隐语的框架下定义同样结构的模型。\n",
    "\n",
    "如前所述，我们需要使用 optimizer 在训练过程中自动地调整模型参数，以使模型在给定的训练数据上达到最佳的性能表现；使用 metric 在机器学习和深度学习中用于衡量模型的性能和表现；使用损失函数监督神经网络的训练。同样地，联邦学习也可以通过使用优化器（optimizer）、衡量指标（metric）和损失函数（loss function） 实现模型更好的性能。在隐语框架下，优化器（optimizer）、衡量指标（metric）和损失函数（loss function）的使用方法和 TensorFlow 中的使用方法一致，也同样通过模型的 compile 函数实现。\n",
    "\n",
    "从迁移过程可以看出，代码修改幅度非常小，整体迁移过程非常方便，充分展现了隐语框架的易用性和便捷性。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_conv_flower_model(input_shape, num_classes, name='model'):\n",
    "    def create_model():\n",
    "        ##########################单机模型代码##################################\n",
    "        from tensorflow import keras\n",
    "\n",
    "        # Create model\n",
    "\n",
    "        model = keras.Sequential(\n",
    "            [\n",
    "                keras.Input(shape=input_shape),\n",
    "                tf.keras.layers.Rescaling(1.0 / 255),\n",
    "                tf.keras.layers.Conv2D(32, 3, activation='relu'),\n",
    "                tf.keras.layers.MaxPooling2D(),\n",
    "                tf.keras.layers.Conv2D(32, 3, activation='relu'),\n",
    "                tf.keras.layers.MaxPooling2D(),\n",
    "                tf.keras.layers.Conv2D(32, 3, activation='relu'),\n",
    "                tf.keras.layers.MaxPooling2D(),\n",
    "                tf.keras.layers.Flatten(),\n",
    "                tf.keras.layers.Dense(128, activation='relu'),\n",
    "                tf.keras.layers.Dense(num_classes),\n",
    "            ]\n",
    "        )\n",
    "        # Compile model\n",
    "        model.compile(\n",
    "            loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
    "            optimizer='adam',\n",
    "            metrics=[\"accuracy\"],\n",
    "        )\n",
    "        return model\n",
    "        ##########################单机模型代码##################################\n",
    "\n",
    "    return create_model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "from secretflow.ml.nn import FLModel\n",
    "from secretflow.security.aggregation import SecureAggregator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:Create proxy actor <class 'secretflow.security.aggregation.secure_aggregator._Masker'> with party alice.\n",
      "INFO:root:Create proxy actor <class 'secretflow.security.aggregation.secure_aggregator._Masker'> with party bob.\n",
      "ERROR:root:Physical devices cannot be modified after being initialized\n",
      "INFO:root:Create proxy actor <class 'secretflow.ml.nn.fl.backend.tensorflow.strategy.fed_avg_w.PYUFedAvgW'> with party alice.\n",
      "INFO:root:Create proxy actor <class 'secretflow.ml.nn.fl.backend.tensorflow.strategy.fed_avg_w.PYUFedAvgW'> with party bob.\n"
     ]
    }
   ],
   "source": [
    "device_list = [alice, bob]\n",
    "aggregator = SecureAggregator(charlie, [alice, bob])\n",
    "\n",
    "# prepare model\n",
    "num_classes = 5\n",
    "input_shape = (180, 180, 3)\n",
    "\n",
    "# keras model\n",
    "model = create_conv_flower_model(input_shape, num_classes)\n",
    "\n",
    "\n",
    "fed_model = FLModel(\n",
    "    device_list=device_list,\n",
    "    model=model,\n",
    "    aggregator=aggregator,\n",
    "    backend=\"tensorflow\",\n",
    "    strategy=\"fed_avg_w\",\n",
    "    random_seed=1234,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 训练和验证模型\n",
    "传入参与方的数据集路径，进行模型的训练和验证"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:FL Train Params: {'x': {PYURuntime(alice): '/tmp/tmp32k3ewwr/flower_photos', PYURuntime(bob): '/tmp/tmp32k3ewwr/flower_photos'}, 'y': None, 'batch_size': 32, 'batch_sampling_rate': None, 'epochs': 5, 'verbose': 1, 'callbacks': None, 'validation_data': {PYURuntime(alice): '/tmp/tmp32k3ewwr/flower_photos', PYURuntime(bob): '/tmp/tmp32k3ewwr/flower_photos'}, 'shuffle': False, 'class_weight': None, 'sample_weight': None, 'validation_freq': 1, 'aggregate_freq': 2, 'label_decoder': None, 'max_batch_size': 20000, 'prefetch_buffer_size': None, 'sampler_method': 'batch', 'random_seed': 1234, 'dp_spent_step_freq': 1, 'audit_log_dir': None, 'dataset_builder': {PYURuntime(alice): <function create_dataset_builder.<locals>.dataset_builder at 0x7fd37cb6cee0>, PYURuntime(bob): <function create_dataset_builder.<locals>.dataset_builder at 0x7fd37cb6d090>}, 'wait_steps': 100, 'self': <secretflow.ml.nn.fl.fl_model.FLModel object at 0x7fd1b9dc07f0>}\n",
      "Train Processing: :   0%|          | 0/31 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/5\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train Processing: :  97%|█████████▋| 30/31 [00:12<00:00,  2.33it/s, {'loss': 0.06979854, 'accuracy': 0.9667013, 'val_loss': 0.0, 'val_accuracy': 1.0}]\n",
      "Train Processing: :   0%|          | 0/31 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 2/5\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train Processing: :  97%|█████████▋| 30/31 [00:11<00:00,  2.52it/s, {'loss': 0.0, 'accuracy': 1.0, 'val_loss': 0.0, 'val_accuracy': 1.0}]\n",
      "Train Processing: :   0%|          | 0/31 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 3/5\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train Processing: :  97%|█████████▋| 30/31 [00:12<00:00,  2.50it/s, {'loss': 0.0, 'accuracy': 1.0, 'val_loss': 0.0, 'val_accuracy': 1.0}]\n",
      "Train Processing: :   0%|          | 0/31 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 4/5\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train Processing: :  97%|█████████▋| 30/31 [00:12<00:00,  2.43it/s, {'loss': 0.0, 'accuracy': 1.0, 'val_loss': 0.0, 'val_accuracy': 1.0}]\n",
      "Train Processing: :   0%|          | 0/31 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 5/5\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train Processing: :  97%|█████████▋| 30/31 [00:12<00:00,  2.35it/s, {'loss': 0.0, 'accuracy': 1.0, 'val_loss': 0.0, 'val_accuracy': 1.0}]\n"
     ]
    }
   ],
   "source": [
    "data = {\n",
    "    alice: path_to_flower_dataset,\n",
    "    bob: path_to_flower_dataset,\n",
    "}\n",
    "history = fed_model.fit(\n",
    "    data,\n",
    "    None,\n",
    "    validation_data=data,\n",
    "    epochs=5,\n",
    "    batch_size=32,\n",
    "    aggregate_freq=2,\n",
    "    sampler_method=\"batch\",\n",
    "    random_seed=1234,\n",
    "    dp_spent_step_freq=1,\n",
    "    dataset_builder=data_builder_dict,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 小结\n",
    "基于 TensorFlow 从单机模型到联邦学习模型，主要需要添加或修改以下几部分\n",
    "- 添加联邦学习中的参与方\n",
    "- 修改数据集的处理逻辑\n",
    "- 进行模型的封装\n",
    "\n",
    "得益于隐语的封装，使用者不需要自己进行大量的代码编写，只需要调用 Secretflow 中的函数，即可便捷完成模型的定义和使用等操作。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 总结\n",
    "本教程说明了使用者能够在 SecretFlow 隐语的框架下，体会到和单机模式下，使用 PyTorch 或 Tensorflow 编程几乎一致的联邦学习模型使用体验。充分展现了隐语框架具有易用性和便捷性等优点。"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
